package dgopenai

import (
	"maps"
	"slices"
	"time"

	dgctx "github.com/darwinOrg/go-common/context"
	dgerr "github.com/darwinOrg/go-common/enums/error"
	"github.com/sashabaranov/go-openai"
)

// ChatGenerateContent generates content for a chat session using a template
func ChatGenerateContent(ctx *dgctx.DgContext, ctgcReq *ChatTemplateGenerateContentRequest) (*GenerateContentResponse, error) {
	templateName := ctgcReq.TemplateName
	templateParams := ctgcReq.TemplateParams
	sessionId := ctgcReq.SessionId

	gcReq, err := MyAigcTemplateProvider.GetGenerateContentRequest(ctx, templateName, templateParams, ctgcReq.MultiContent)
	if err != nil {
		return nil, err
	}
	if ctgcReq.RequestId != "" {
		gcReq.RequestId = ctgcReq.RequestId
	}

	historyMessages, err := MyChatMessageProvider.FindChatMessages(ctx, sessionId)
	if err != nil {
		return nil, err
	}

	if len(historyMessages) > 0 {
		var messages []openai.ChatCompletionMessage
		for _, historyMessage := range historyMessages {
			messages = append(messages, buildChatCompletionMessageByChatMessage(historyMessage))
		}
		if len(gcReq.Messages) > 0 {
			messages = append(messages, gcReq.Messages...)
		}
		gcReq.Messages = messages
	}

	fixGenerateContentRequest(gcReq)

	var (
		gcResp *GenerateContentResponse
		retry  int
		cost   time.Duration
	)

	defer func() {
		go func() {
			chatMessage := buildChatMessageByGenerateContentResponse(gcResp)
			_ = MyChatMessageProvider.SaveChatMessages(ctx, sessionId, []*ChatMessage{chatMessage})
			_ = MyAigcResponseProvider.SaveChatGenerateContentResponse(ctx, templateName, gcReq, gcResp, ctgcReq.BizParams, sessionId, chatMessage.Id, retry, cost, err)
		}()
	}()

	for i := 0; i < gcReq.TryTimes; i++ {
		for j, platformModel := range gcReq.PlatformModels {
			retry = i*len(gcReq.PlatformModels) + j
			gcResp, cost, err = createChatSessionCompletion(ctx, chatAigcKey, templateName, platformModel, gcReq, nil, sessionId)
			if err != nil {
				continue
			}

			return gcResp, nil
		}
	}

	return nil, dgerr.SYSTEM_ERROR
}

// Chat handles a chat request
func Chat(ctx *dgctx.DgContext, chatReq *ChatRequest) (*ChatResponse, error) {
	templateName, gcReq, err := buildChatGenerateContentRequest(ctx, chatReq)
	if err != nil {
		return nil, err
	}
	sessionId := chatReq.SessionId

	var (
		gcResp   *GenerateContentResponse
		chatResp *ChatResponse
		retry    int
		cost     time.Duration
	)

	defer func() {
		if chatResp != nil {
			_ = MyChatMessageProvider.SaveChatMessages(ctx, sessionId, chatResp.Messages)
			messageId := chatResp.Messages[len(chatResp.Messages)-1].Id
			_ = MyAigcResponseProvider.SaveChatGenerateContentResponse(ctx, templateName, gcReq, gcResp, chatReq.Params, sessionId, messageId, retry, cost, err)
		}
	}()

	for i := 0; i < gcReq.TryTimes; i++ {
		for j, platformModel := range gcReq.PlatformModels {
			retry = i*len(gcReq.PlatformModels) + j
			gcResp, cost, err = createChatSessionCompletion(ctx, chatKey, templateName, platformModel, gcReq, chatReq.Params, sessionId)
			if gcResp == nil {
				continue
			}

			chatResp = buildChatResponse(chatReq, gcResp)
			return chatResp, nil
		}
	}

	return nil, dgerr.SYSTEM_ERROR
}

// StreamChat handles a chat request and streams the response
func StreamChat(ctx *dgctx.DgContext, chatReq *ChatRequest) (chan *GenerateContentResponse, error) {
	templateName, gcReq, err := buildChatGenerateContentRequest(ctx, chatReq)
	if err != nil {
		return nil, err
	}
	sessionId := chatReq.SessionId

	var (
		mergedResp *GenerateContentResponse
		retry      int
		cost       time.Duration
	)

	for i := 0; i < gcReq.TryTimes; i++ {
		for j, platformModel := range gcReq.PlatformModels {
			retry = i*len(gcReq.PlatformModels) + j
			var (
				stream    *openai.ChatCompletionStream
				bizParams map[string]any
				logParams []any
			)
			stream, bizParams, logParams, err = createChatSessionCompletionStream(ctx, chatKey, templateName, platformModel, gcReq, chatReq.Params, sessionId)
			if err != nil {
				continue
			}

			respChan := make(chan *GenerateContentResponse, 10)
			go func() {
				defer func() {
					_ = stream.Close()
					close(respChan)

					if mergedResp != nil {
						chatResp := buildChatResponse(chatReq, mergedResp)
						_ = MyChatMessageProvider.SaveChatMessages(ctx, sessionId, chatResp.Messages)
						messageId := chatResp.Messages[len(chatResp.Messages)-1].Id
						_ = MyAigcResponseProvider.SaveChatGenerateContentResponse(ctx, templateName, gcReq, mergedResp, chatReq.Params, sessionId, messageId, retry, cost, err)
					} else {
						_ = MyAigcResponseProvider.SaveChatGenerateContentResponse(ctx, templateName, gcReq, mergedResp, chatReq.Params, sessionId, 0, retry, cost, err)
					}
				}()

				mergedResp, cost, err = receiveChatCompletionStream(ctx, chatKey, stream, respChan, platformModel, gcReq, bizParams, logParams)
			}()

			return respChan, nil
		}
	}

	return nil, dgerr.SYSTEM_ERROR
}

// createChatSessionCompletion creates a chat completion for a session
func createChatSessionCompletion(ctx *dgctx.DgContext, gcType string, templateName string, platformModel *PlatformModelPair, gcReq *GenerateContentRequest, params map[string]any, sessionId int64) (*GenerateContentResponse, time.Duration, error) {
	bizParams := map[string]any{sessionIdKey: sessionId}
	if len(params) > 0 {
		maps.Copy(bizParams, params)
	}

	return createChatCompletion(ctx, gcType, templateName, platformModel, gcReq, bizParams)
}

// createChatSessionCompletionStream creates a chat completion stream for a session
func createChatSessionCompletionStream(ctx *dgctx.DgContext, gcType string, templateName string, platformModel *PlatformModelPair, gcReq *GenerateContentRequest, params map[string]any, sessionId int64) (*openai.ChatCompletionStream, map[string]any, []any, error) {
	bizParams := map[string]any{sessionIdKey: sessionId}
	if len(params) > 0 {
		maps.Copy(bizParams, params)
	}

	stream, logParams, err := createChatCompletionStream(ctx, gcType, templateName, platformModel, gcReq, bizParams)
	return stream, bizParams, logParams, err
}

// buildChatGenerateContentRequest builds a GenerateContentRequest for a chat request
func buildChatGenerateContentRequest(ctx *dgctx.DgContext, chatReq *ChatRequest) (string, *GenerateContentRequest, error) {
	sessionId := chatReq.SessionId
	templateName, err := MyChatSessionProvider.GetTemplateNameBySessionId(ctx, sessionId)
	if err != nil {
		return "", nil, err
	}

	gcReq, err := MyAigcTemplateProvider.GetChatGenerateContentRequest(ctx, templateName)
	if err != nil {
		return "", nil, err
	}
	if chatReq.RequestId != "" {
		gcReq.RequestId = chatReq.RequestId
	}

	historyMessages, err := MyChatMessageProvider.FindChatMessages(ctx, sessionId)
	if err != nil {
		return "", nil, err
	}

	// 如果启用了思考模式，则删除历史记录中第一个Assistant消息
	if gcReq.EnableThinking && len(historyMessages) > 0 {
		var (
			isDeleteFirstAssistant bool
			firstAssistantIndex    int
		)
		for i, historyMessage := range historyMessages {
			if historyMessage.Role == openai.ChatMessageRoleUser {
				break
			} else if historyMessage.Role == openai.ChatMessageRoleAssistant {
				isDeleteFirstAssistant = true
				firstAssistantIndex = i
			}
		}
		if isDeleteFirstAssistant {
			historyMessages = slices.Delete(historyMessages, firstAssistantIndex, firstAssistantIndex+1)
		}
	}

	if len(historyMessages) > 0 {
		var messages []openai.ChatCompletionMessage
		for _, historyMessage := range historyMessages {
			messages = append(messages, buildChatCompletionMessageByChatMessage(historyMessage))
		}
		currMessage := openai.ChatCompletionMessage{
			Role:         openai.ChatMessageRoleUser,
			Content:      chatReq.Content,
			MultiContent: chatReq.MultiContent,
		}
		if len(gcReq.Messages) > 0 {
			currMessage.FunctionCall = gcReq.Messages[len(gcReq.Messages)-1].FunctionCall
		}
		messages = append(messages, currMessage)
		gcReq.Messages = messages
	}

	fixGenerateContentRequest(gcReq)
	return templateName, gcReq, nil
}

// buildChatResponse builds a ChatResponse from a ChatRequest and GenerateContentResponse
func buildChatResponse(chatReq *ChatRequest, gcResp *GenerateContentResponse) *ChatResponse {
	chatMessage := buildChatMessageByChatRequest(chatReq)
	if chatMessage.RequestId == "" {
		chatMessage.RequestId = gcResp.RequestId
	}
	respMessage := buildChatMessageByGenerateContentResponse(gcResp)
	return &ChatResponse{
		Messages: []*ChatMessage{chatMessage, respMessage},
	}
}
